#!/bin/bash
{
	#////////////////////////////////////
	# DietPi Function:
	# - benchmark
	#
	#////////////////////////////////////
	# Created by Daniel Knight / daniel.knight@dietpi.com / dietpi.com
	#
	#////////////////////////////////////
	#
	# Location: /boot/dietpi/func/dietpi-benchmark
	#
	# Usage:
	# /boot/dietpi/func/dietpi-benchmark 0								= init vars for sourcing /var/lib/dietpi/dietpi-benchmark/results
	# FP_BENCHFILE=/location BENCH_FILESIZE=optional_size_MiB /boot/dietpi/func/dietpi-benchmark 1	= Benchmark $FP_BENCHFILE filesystem read/write.
	# /boot/dietpi/func/dietpi-benchmark 2								= Run all benchmarks and upload to dietpi.com
	# G_INTERACTIVE=0 /boot/dietpi/func/dietpi-benchmark 2						= Same as above, automated
	# /boot/dietpi/func/dietpi-benchmark 3								= iPerf server
	# /boot/dietpi/func/dietpi-benchmark 4								= iPerf client with prompt
	#
	#////////////////////////////////////

	# Import DietPi-Globals --------------------------------------------------------------
	. /boot/dietpi/func/dietpi-globals
	readonly G_PROGRAM_NAME='DietPi-Benchmark'
	G_INIT
	# Import DietPi-Globals --------------------------------------------------------------

	INPUT=$1
	if [[ ! $INPUT ]]; then

		G_DIETPI-NOTIFY 1 'Invalid command, please see sourcecode'
		exit 1

	fi

	#/////////////////////////////////////////////////////////////////////////////////////
	# Global
	#/////////////////////////////////////////////////////////////////////////////////////
	SURVEY_OPTED_IN=0

	IPERF_SERVER_IP=
	SHOW_RESULTS=${SHOW_RESULTS:-1} # Optional input: Show results for individual tests?
	CPUBENCH_INT_MAX=${CPUBENCH_INT_MAX:-1000000} # Optional input
	BENCH_FILESIZE=${BENCH_FILESIZE:-0} # Optional input for automated
	[[ ${FP_BENCHFILE:=/} == '/'* ]] || { G_DIETPI-NOTIFY 1 'FP_BENCHFILE must be an absolute file path, aborting...'; exit 1; } # Optional input

	BENCH_ROOTFS_WRITE='Not tested'
	BENCH_ROOTFS_READ='Not tested'
	BENCH_RAM_WRITE='Not tested'
	BENCH_RAM_READ='Not tested'
	BENCH_CUSTOMFS_WRITE='Not tested'
	BENCH_CUSTOMFS_READ='Not tested'
	BENCH_CPU='Not tested'
	BENCH_CPU_TEMP_START='Not tested'
	BENCH_CPU_TEMP_END='Not tested'
	BENCH_NET_LAN_SPEED='Not tested'

	Load_Results(){ [[ -f '/var/lib/dietpi/dietpi-benchmark/results' ]] && . /var/lib/dietpi/dietpi-benchmark/results; }

	Save_Results(){

		G_EXEC mkdir -p /var/lib/dietpi/dietpi-benchmark
		cat << _EOF_ > /var/lib/dietpi/dietpi-benchmark/results
# -------------------------
# DietPi-Benchmark
# -------------------------
BENCH_VERSION=2
BENCH_HW_MODEL=$G_HW_MODEL
BENCH_CPU='$BENCH_CPU'
BENCH_ROOTFS_WRITE='$BENCH_ROOTFS_WRITE'
BENCH_ROOTFS_READ='$BENCH_ROOTFS_READ'
BENCH_RAM_WRITE='$BENCH_RAM_WRITE'
BENCH_RAM_READ='$BENCH_RAM_READ'
BENCH_CUSTOMFS_WRITE='$BENCH_CUSTOMFS_WRITE'
BENCH_CUSTOMFS_READ='$BENCH_CUSTOMFS_READ'
BENCH_CPU_TEMP_START='$BENCH_CPU_TEMP_START'
BENCH_CPU_TEMP_END='$BENCH_CPU_TEMP_END'
BENCH_NET_LAN_SPEED='$BENCH_NET_LAN_SPEED'
_EOF_
	}

	Update_Survey_Opted_Status(){ SURVEY_OPTED_IN=$(sed -n '/^[[:blank:]]*SURVEY_OPTED_IN=[01]$/{s/^[^=]*=//p;q}' /boot/dietpi.txt); }

	Upload_Survey(){

		Update_Survey_Opted_Status

		if [[ $SURVEY_OPTED_IN == 1 ]]; then

			/boot/dietpi/dietpi-survey 1

		elif G_WHIP_YESNO 'To upload and compare your results against other users, you must be opted in to the DietPi-Survey.\n\nWould you like to run DietPi-Survey now, allowing you to opt in?'; then

			/boot/dietpi/dietpi-survey

		fi

	}

	CPU_Benchmark(){

		G_DIETPI-NOTIFY 0 "Running DietPi-CPU Benchmark with max value of 1000000 over $G_HW_CPU_CORES cores, please wait..."

		cat << _EOF_ > bench
#!/bin/bash
target_max_int=$CPUBENCH_INT_MAX
cores=$G_HW_CPU_CORES
int_split=\$(( \$target_max_int / \$cores ))
aStart_Int=()
aEnd_Int=()

# Split the max int target based on total cores
for ((i=0; i<\$cores; i++))
do

	aEnd_Int[\$i]=\$(( (\$i + 1) * \$int_split ))
	aStart_Int[\$i]=\$(( \${aEnd_Int[\$i]} - \$int_split ))
	echo \${aStart_Int[\$i]} \${aEnd_Int[\$i]}

done

Run_Bench(){

	while (( \${aStart_Int[\$1]} < \${aEnd_Int[\$1]} ))
	do

		((aStart_Int[\$1]++))

	done

}

# Launch benchmark threads
for ((i=0; i<\$cores; i++))
do

	Run_Bench \$i &

done

# Wait for jobs to finish
for job in \`jobs -p\`
do

	echo \$job
	wait \$job

done
_EOF_
		G_EXEC chmod +x bench

		sync
		sleep 1

		BENCH_CPU_TEMP_START=$(G_OBTAIN_CPU_TEMP)

		{ time -p ./bench; } 2> bench_result

		BENCH_CPU_TEMP_END=$(G_OBTAIN_CPU_TEMP)
		BENCH_CPU=$(mawk '/real/{print $2;exit}' bench_result)

		if (( $SHOW_RESULTS )); then

			G_WHIP_MSG "CPU Benchmark Results:\n
 - Total time     = $BENCH_CPU Seconds
 - Value to reach = $CPUBENCH_INT_MAX
 - CPU temps      = Start: ${BENCH_CPU_TEMP_START}'c | End: ${BENCH_CPU_TEMP_END}'c
\n\nNB: A lower 'Total time' is faster"

		fi

	}

	Filesystem_Benchmark(){

		G_EXEC mkdir -p "$FP_BENCHFILE" # Failsafe
		local min_benchmark_size=10
		local max_benchmark_size=$(( $(G_CHECK_FREESPACE "$FP_BENCHFILE") - 50 )) # Assure 50 MiB stays free
		#(( $max_benchmark_size > 10000 )) && max_benchmark_size=10000

		local benchmark_file_name='benchmark.file'
		FP_BENCHFILE="${FP_BENCHFILE%/}/$benchmark_file_name"

		G_TERM_CLEAR

		# Obtain user benchmark size
		if (( $BENCH_FILESIZE == 0 )); then

			G_WHIP_DEFAULT_ITEM=$min_benchmark_size
			G_WHIP_INPUTBOX "Please enter a filesize test value (MiB). A value of 100 is recommended.
 - Minimum value: $min_benchmark_size
 - Maximum possible value: $max_benchmark_size" && BENCH_FILESIZE=$G_WHIP_RETURNED_VALUE

		fi

		# Sanity
		disable_error=1 G_CHECK_VALIDINT "$BENCH_FILESIZE" $min_benchmark_size || BENCH_FILESIZE=$min_benchmark_size
		(( $BENCH_FILESIZE > $max_benchmark_size )) && BENCH_FILESIZE=$max_benchmark_size

		# Run benchmark
		G_DIETPI-NOTIFY 2 "Performing Read/Write benchmark on $FP_BENCHFILE, please wait..."

		G_DIETPI-NOTIFY 2 "Testing Seq Write Speed ($BENCH_FILESIZE MiB)"
		sleep 1
		local bench_write_result=$(dd bs=4K count=$(( $BENCH_FILESIZE * 256 )) if=/dev/zero of="$FP_BENCHFILE" conv=fdatasync 2>&1 | grep -m1 'copied')
		bench_write_result=$(mawk '{printf "%.0f",$1/$8/1024^2}' <<< "$bench_write_result")

		# - Clear cache
		sync
		echo 3 > /proc/sys/vm/drop_caches

		G_DIETPI-NOTIFY 2 "Testing Seq Read Speed ($BENCH_FILESIZE MiB)"
		sleep 1
		local bench_read_result=$(dd bs=4K count=$(( $BENCH_FILESIZE * 256 )) if="$FP_BENCHFILE" of=/dev/null 2>&1 | grep -m1 'copied')
		bench_read_result=$(mawk '{printf "%.0f",$1/$8/1024^2}' <<< "$bench_read_result")

		# - RootFS
		if [[ $FP_BENCHFILE == '/'$benchmark_file_name ]]; then

			BENCH_ROOTFS_WRITE=$bench_write_result
			BENCH_ROOTFS_READ=$bench_read_result

		# - RAM
		elif [[ $FP_BENCHFILE == '/tmp/'$benchmark_file_name ]]; then

			BENCH_RAM_WRITE=$bench_write_result
			BENCH_RAM_READ=$bench_read_result

		# - Custom
		else

			BENCH_CUSTOMFS_WRITE=$bench_write_result
			BENCH_CUSTOMFS_READ=$bench_read_result

		fi

		if (( $SHOW_RESULTS )); then

			G_WHIP_MSG "Filesystem Benchmark Results:\n
 - Filepath  = $FP_BENCHFILE
 - Test size = $BENCH_FILESIZE MiB
 - WRITE     = $bench_write_result MiB/s
 - READ      = $bench_read_result MiB/s"

		fi

		[[ -f $FP_BENCHFILE ]] && rm -v "$FP_BENCHFILE"

	}

	Benchmark_Survey(){

		SHOW_RESULTS=0
		G_EXEC mkdir -p /var/lib/dietpi/dietpi-benchmark

		if (( $G_INTERACTIVE )); then

			if G_WHIP_YESNO "$G_PROGRAM_NAME will now run the following benchmarks:\n - CPU performance\n - RootFS read/write performance\n - RAM read/write performance.
The test duration will depend on hardware specs, however, most devices will not exceed 40 seconds.\n\nIf you are opted in to the DietPi-Survey, your scores will uploaded automatically.
You can compare results against other users and hardware using the following link:\n - https://dietpi.com/survey#benchmark\n\nPlease select 'Ok' to begin the benchmark tests."; then

				# Additional checks
				(( $G_HW_CPU_CORES > 2 )) || G_WHIP_YESNO "We have detected a low CPU core count ($G_HW_CPU_CORES). The tests may take up to 5 minutes to complete.\n\nDo you wish to continue?" || exit 0

			else

				G_DIETPI-NOTIFY 0 'Benchmark aborted'
				exit 0

			fi

		fi

		G_DIETPI-NOTIFY 0 'Running benchmark suite tests, please wait...'
		CPU_Benchmark

		# FS/RAM
		# - RAM
		if df -t tmpfs /tmp &> /dev/null; then

			FP_BENCHFILE='/tmp'
			BENCH_FILESIZE=$(mawk '/^MemFree:/{printf "%.0f",$2/2048;exit}' /proc/meminfo) # Assure 50% of physical mem stays free, so swapping can be ruled out: KiB => MiB / 2
			Filesystem_Benchmark

		else

			G_DIETPI-NOTIFY 1 '/tmp does not seem to be mounted as tmpfs (RAMdisk). Skipping RAM benchmark...'

		fi

		# - RootFS
		FP_BENCHFILE='/'
		BENCH_FILESIZE=100
		Filesystem_Benchmark

		SHOW_RESULTS=1

		G_WHIP_MSG "Benchmarks completed:
 - CPU Performance : Duration = $BENCH_CPU seconds (lower is faster)
 - CPU Temp        : Idle = ${BENCH_CPU_TEMP_START}'c | Full load = ${BENCH_CPU_TEMP_END}'c
 - RootFS          : Write = $BENCH_ROOTFS_WRITE MiB/s | Read = $BENCH_ROOTFS_READ MiB/s
 - RAM             : Write = $BENCH_RAM_WRITE MiB/s | Read = $BENCH_RAM_READ MiB/s

Compare these results online with other users, using the link below:
 - https://dietpi.com/survey#benchmark"

	}

	Benchmark_Lan_Server(){

		G_AG_CHECK_INSTALL_PREREQ iperf
		G_THREAD_START iperf -s
		G_WHIP_MSG "Network benchmark is now running in 'server mode'. Please use another DietPi system to connect as a client and begin the test.
 - Server IP = $(mawk 'NR==4' /run/dietpi/.network)\n\nWhen the client has finished, select 'Ok' to stop the server."
		killall -w iperf
		G_THREAD_WAIT

	}

	Benchmark_Lan_Client(){

		G_AG_CHECK_INSTALL_PREREQ iperf
		if G_WHIP_INPUTBOX 'Please enter the IP address of the system running the network benchmark server:\n - eg: 192.168.0.100'; then

			IPERF_SERVER_IP=$G_WHIP_RETURNED_VALUE

			G_DIETPI-NOTIFY 2 "Benchmarking network performance to $IPERF_SERVER_IP. Please wait..."
			BENCH_NET_LAN_SPEED=$(iperf -c "$IPERF_SERVER_IP" -p 5001 -t 5 -y C | sed 's/.*,//g')
			if disable_error=1 G_CHECK_VALIDINT "$BENCH_NET_LAN_SPEED"; then

				BENCH_NET_LAN_SPEED=$(( $BENCH_NET_LAN_SPEED / 8 / 1024**2 ))
				G_WHIP_MSG "Network LAN benchmark completed:\n - Transfer rate = $BENCH_NET_LAN_SPEED MiB/s"

			else

				G_WHIP_MSG 'Network LAN benchmark failed, please verify:\n - Server is running on the other system\n - Server IP address is correct'
				BENCH_NET_LAN_SPEED='Not tested'

			fi

		fi

	}

	#/////////////////////////////////////////////////////////////////////////////////////
	# Main Loop
	#/////////////////////////////////////////////////////////////////////////////////////
	#-----------------------------------------------------------------------------------
	# Run at highest CPU and I/O priority
	renice -n -19 $$
	ionice -p $$ -c 1 -n 0
	#-----------------------------------------------------------------------------------
	# Init vars for sourcing
	if (( $INPUT == 0 )); then

		Save_Results

	# LAN bench server
	elif (( $INPUT == 3 )); then

		Benchmark_Lan_Server

	# Run benchmarks
	else

		#-----------------------------------------------------------------------------------
		/boot/dietpi/dietpi-services stop
		#-----------------------------------------------------------------------------------
		# Reuse existing results if done
		Load_Results
		#-----------------------------------------------------------------------------------
		if (( $INPUT == 1 )); then

			Filesystem_Benchmark

		elif (( $INPUT == 2 )); then

			Benchmark_Survey

		elif (( $INPUT == 4 )); then

			Benchmark_Lan_Client

		fi
		#-----------------------------------------------------------------------------------
		# Export results data for other scripts:
		Save_Results
		#-----------------------------------------------------------------------------------
		# Upload if in survey, else, prompt
		Upload_Survey
		#-----------------------------------------------------------------------------------
		/boot/dietpi/dietpi-services start
		#-----------------------------------------------------------------------------------

	fi
	#-----------------------------------------------------------------------------------
	exit 0
	#-----------------------------------------------------------------------------------
}
